package test

import (
	"encoding/json"
	"log"
	"net/http"

	"github.com/gorilla/websocket"
	"github.com/pion/webrtc/v3"
)

var (
	upgrader = websocket.Upgrader{}
	peers    = map[*websocket.Conn]*webrtc.PeerConnection{}
)

func main() {
	http.HandleFunc("/offer", handleOffer)
	http.HandleFunc("/answer", handleAnswer)
	http.HandleFunc("/ws", handleWebSocket)
	log.Fatal(http.ListenAndServe(":8080", nil))
}

func handleOffer(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	pc, err := webrtc.NewPeerConnection(webrtc.Configuration{})
	if err != nil {
		log.Println(err)
		return
	}

	pc.OnICECandidate(func(candidate *webrtc.ICECandidate) {
		if candidate != nil {
			jsonCandidate, _ := json.Marshal(candidate.ToJSON())
			conn.WriteMessage(websocket.TextMessage, jsonCandidate)
		}
	})

	pc.OnTrack(func(track *webrtc.TrackRemote, receiver *webrtc.RTPReceiver) {
		go func() {
			for {
				_, _, _ = track.ReadRTP()
			}
		}()
	})

	//pc.AddTransceiver(webrtc.RTPCodecTypeVideo)
	pc.AddTransceiverFromKind(webrtc.RTPCodecTypeVideo)
	peers[conn] = pc

	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			log.Println(err)
			delete(peers, conn)
			return
		}
		pc.AddICECandidate(webrtc.ICECandidateInit{
			Candidate: string(message),
		})
	}
}

// handleAnswer 处理answer请求
func handleAnswer(w http.ResponseWriter, r *http.Request) {
	var answer struct {
		Candidate string `json:"candidate"`
	}
	err := json.NewDecoder(r.Body).Decode(&answer)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	pc := peers[conn]
	pc.AddICECandidate(webrtc.ICECandidateInit{
		Candidate: answer.Candidate,
	})
}

// handleWebSocket 处理WebSocket连接请求
func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	pc, err := webrtc.NewPeerConnection(webrtc.Configuration{})
	if err != nil {
		log.Println(err)
		return
	}

	pc.OnICECandidate(func(candidate *webrtc.ICECandidate) {
		if candidate != nil {
			jsonCandidate, _ := json.Marshal(candidate.ToJSON())
			conn.WriteMessage(websocket.TextMessage, jsonCandidate)
		}
	})

	pc.OnTrack(func(track *webrtc.TrackRemote, receiver *webrtc.RTPReceiver) {
		go func() {
			for {
				_, _, _ = track.ReadRTP()
			}
		}()
	})

	offer, err := pc.CreateOffer(nil)
	if err != nil {
		log.Println(err)
		return
	}

	err = pc.SetLocalDescription(offer)
	if err != nil {
		log.Println(err)
		return
	}

	conn.WriteJSON(offer)

	peers[conn] = pc
}
